/*
 * This file is part of Dependency-Track.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *
 * SPDX-License-Identifier: Apache-2.0
 * Copyright (c) Steve Springett. All Rights Reserved.
 */
package org.dependencytrack.util;

import org.dependencytrack.model.Severity;
import java.math.BigDecimal;

public final class VulnerabilityUtil {

    private VulnerabilityUtil() { }

    /**
     * Returns the severity based on the numerical CVSS score. CVSSv2 and CVSSv3 have
     * slightly different ranges with CVSSv3 introducing critical severity whereas
     * CVSSv2 only has high, medium, and low.
     *
     * This method properly accounts for vulnerabilities that may have only a CVSSv2
     * score. If both scores are available, it will return the CVSSv3 severity.
     * @return the severity of the vulnerability
     * @since 3.1.0
     */
    public static Severity getSeverity(final Object cvssV2BaseScore, final Object cvssV3BaseScore) {
        if (cvssV3BaseScore instanceof BigDecimal) {
            return normalizedCvssV3Score(((BigDecimal)cvssV3BaseScore).doubleValue());
        } else if (cvssV2BaseScore instanceof BigDecimal) {
            return normalizedCvssV2Score(((BigDecimal)cvssV2BaseScore).doubleValue());
        } else {
            return Severity.UNASSIGNED;
        }
    }

    /**
     * Returns the severity based on the numerical CVSS score. CVSSv2 and CVSSv3 have
     * slightly different ranges with CVSSv3 introducing critical severity whereas
     * CVSSv2 only has high, medium, and low.
     *
     * This method properly accounts for vulnerabilities that may have only a CVSSv2
     * score. If both scores are available, it will return the CVSSv3 severity.
     * @return the severity of the vulnerability
     * @since 3.2.1
     */
    public static Severity getSeverity(final Object severity, final Object cvssV2BaseScore, final Object cvssV3BaseScore) {
        if (severity instanceof String) {
            final String s = (String)severity;
            if (s.equalsIgnoreCase(Severity.CRITICAL.name())) {
                return Severity.CRITICAL;
            } else if (s.equalsIgnoreCase(Severity.HIGH.name())) {
                return Severity.HIGH;
            } else if (s.equalsIgnoreCase(Severity.MEDIUM.name())) {
                return Severity.MEDIUM;
            } else if (s.equalsIgnoreCase(Severity.LOW.name())) {
                return Severity.LOW;
            } else if (s.equalsIgnoreCase(Severity.INFO.name())) {
                return Severity.INFO;
            }
        } else if (severity instanceof Severity) {
            return (Severity)severity;
        } else {
            return getSeverity(cvssV2BaseScore, cvssV3BaseScore);
        }
        return Severity.UNASSIGNED;
    }

    /**
     * Returns the severity based on the numerical CVSS score. CVSSv2 and CVSSv3 have
     * slightly different ranges with CVSSv3 introducing critical severity whereas
     * CVSSv2 only has high, medium, and low.
     *
     * This method properly accounts for vulnerabilities that may have only a CVSSv2
     * score. If both scores are available, it will return the CVSSv3 severity.
     * @return the severity of the vulnerability
     * @since 3.1.0
     */
    public static Severity getSeverity(final BigDecimal cvssV2BaseScore, final BigDecimal cvssV3BaseScore) {
        if (cvssV3BaseScore != null) {
            return normalizedCvssV3Score(cvssV3BaseScore.doubleValue());
        } else if (cvssV2BaseScore != null) {
            return normalizedCvssV2Score(cvssV2BaseScore.doubleValue());
        } else {
            return Severity.UNASSIGNED;
        }
    }

    /**
     * Returns the severity based on the numerical CVSS score.
     * @return the severity of the vulnerability
     * @since 3.1.0
     */
    public static Severity normalizedCvssV2Score(final double score) {
        if (score >= 7) {
            return Severity.HIGH;
        } else if (score >= 4) {
            return Severity.MEDIUM;
        } else if (score > 0) {
            return Severity.LOW;
        } else {
            return Severity.UNASSIGNED;
        }
    }

    /**
     * Returns the severity based on the numerical CVSS score.
     * @return the severity of the vulnerability
     * @since 3.1.0
     */
    public static Severity normalizedCvssV3Score(final double score) {
        if (score >= 9) {
            return Severity.CRITICAL;
        } else if (score >= 7) {
            return Severity.HIGH;
        } else if (score >= 4) {
            return Severity.MEDIUM;
        } else if (score > 0) {
            return Severity.LOW;
        } else {
            return Severity.UNASSIGNED;
        }
    }

}
